package de.yvert.jingle.impl.openexr;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.zip.DataFormatException;
import java.util.zip.Inflater;

import de.yvert.jingle.hdr.HdrImage2D;
import de.yvert.jingle.help.Buffer;
import de.yvert.jingle.help.ConverterInputStream;

public class ExrParser
{

private static final int MAGIC = 20000630; // 30.06.2000

private static final int FLAG_TILES = 2;

ConverterInputStream in;
HashMap<String,Attribute> attributes = new HashMap<String,Attribute>();

public ExrParser(InputStream in)
{
	this.in = new ConverterInputStream(ByteOrder.LITTLE_ENDIAN, in);
}

private String readNullString() throws IOException
{
	int value = in.readUnsignedByte();
	if (value == 0) return null;
	
	byte[] data = new byte[32];
	int index = 0;
	while (value != 0)
	{
		data[index++] = (byte) value;
		value = in.readUnsignedByte();
	}
	return new String(data, 0, index);
}

private Attribute readChList(String name, String type, int size) throws IOException
{
	if (!ChannelListAttribute.TYPE.equals(type)) throw new IOException("Argh!");
	ArrayList<Channel> channels = new ArrayList<Channel>();
	String chname;
	while ((chname = readNullString()) != null)
	{
		Channel c = new Channel();
		c.name = chname;
		c.pixelType = PixelType.decode(in.readInt());
		c.reserved = in.readInt();
		c.xSampling = in.readInt();
		c.ySampling = in.readInt();
		channels.add(c);
	}
	ChannelListAttribute result = new ChannelListAttribute(name, new ChannelList(channels.toArray(new Channel[0])));
	if (size != result.calculateSize()) throw new IOException("Argh!");
	return result;
}

private Attribute readBox2i(String name, String type, int size) throws IOException
{
	if (!Box2iAttribute.TYPE.equals(type)) throw new IOException("Argh!");
	if (Box2iAttribute.SIZE != size) throw new IOException("Argh!");
	int x0 = in.readInt();
	int y0 = in.readInt();
	int x1 = in.readInt();
	int y1 = in.readInt();
	return new Box2iAttribute(name, x0, y0, x1, y1);
}

private Attribute readCompression(String name, String type, int size) throws IOException
{
	if (!CompressionAttribute.TYPE.equals(type)) throw new IOException("Argh!");
	if (CompressionAttribute.SIZE != size) throw new IOException("Argh!");
	int c = in.readUnsignedByte();
	CompressionType ctype = CompressionType.NO_COMPRESSION;
	switch (c)
	{
		case 0 : ctype = CompressionType.NO_COMPRESSION; break;
		case 1 : ctype = CompressionType.RLE_COMPRESSION; break;
		case 2 : ctype = CompressionType.ZIPS_COMPRESSION; break;
		case 3 : ctype = CompressionType.ZIP_COMPRESSION; break;
		case 4 : ctype = CompressionType.PIZ_COMPRESSION; break;
		case 5 : ctype = CompressionType.PXR24_COMPRESSION; break;
		default :
			throw new IOException("Unsupported Compression Type: "+c);
	}
	return new CompressionAttribute(StdAttributes.COMPRESSION, ctype);
}

private Attribute readAttribute() throws IOException
{
	String name = readNullString();
	if (name == null) return null;
	String type = readNullString();
	int size = in.readInt();
	
	if (ChannelListAttribute.TYPE.equals(type))
		return readChList(name, type, size);
	if (Box2iAttribute.TYPE.equals(type))
		return readBox2i(name, type, size);
	if (CompressionAttribute.TYPE.equals(type))
		return readCompression(name, type, size);
	
	GenericAttribute result = new GenericAttribute(name);
	result.type = type;
	byte[] data = new byte[size];
	for (int i = 0; i < size; i++)
		data[i] = in.readByte();
	result.value = data;
	return result;
}

public HdrImage2D parse() throws IOException
{
	int m = in.readInt();
	if (m != MAGIC) throw new IOException("File is not an OpenEXR file!");
	int vflags = in.readInt();
	int version = vflags & 0xff;
	if (version != 2) throw new IOException("Can only read version 2 files!");
	int flags = vflags >> 8;
	if ((flags & FLAG_TILES) != 0) throw new IOException("Cannot read tile-based files!");
	if ((flags & ~FLAG_TILES) != 0) throw new IOException("Unknown flags set - bailing out!");
	
	System.out.println("Header:");
	Attribute attr;
	while ((attr = readAttribute()) != null)
	{
		attributes.put(attr.getName(), attr);
		if (!"preview".equals(attr.getName()))
			System.out.println("  "+attr);
	}
	
	ChannelList channels = ((ChannelListAttribute) attributes.get(StdAttributes.CHANNELS)).getChannels();
	if ((channels.size() != 3) && (channels.size() != 4) && (channels.size() != 5)) throw new IOException("Not implemented!");
	int numChannels = channels.size();
	
	PixelType pixelType = channels.get(0).getPixelType();
	for (int i = 0; i < numChannels; i++)
	{
		Channel c = channels.get(i);
		if (c.getPixelType() != pixelType)
			throw new IOException("Not implemented!");
		if ((c.getXSampling() != 1) || (c.getYSampling() != 1))
			throw new IOException("Not implemented!");
	}
	if ((pixelType != PixelType.HALF) && (pixelType != PixelType.FLOAT))
		throw new IOException("Not implemented!");
	
	Box2iAttribute dw = (Box2iAttribute) attributes.get(StdAttributes.DATA_WINDOW);
	int height = dw.getHeight(), width = dw.getWidth();
	
	CompressionType compression = ((CompressionAttribute) attributes.get(StdAttributes.COMPRESSION)).getType();
	
	int blockSize = compression.getBlockSize();
	int numBlocks = (height+blockSize-1)/blockSize;
	for (int i = 0; i < numBlocks; i++)
	{
		/*long offset = */ in.readLong();
//		System.out.println("  Offset: "+offset);
	}
	
	final int maxBlockSize = blockSize*channels.getLineLength(width);
	
	Decompressor decompressor;
	if (compression == CompressionType.ZIP_COMPRESSION)
		decompressor = new Decompressor()
			{
				private final byte[] decompressed = new byte[maxBlockSize];
				private final Inflater inflater = new Inflater();
				@Override
				public int decompress(byte[] compressed, int size, byte[] pixeldata) throws IOException
				{
					inflater.reset();
					inflater.setInput(compressed, 0, size);
					int resultLength;
					try
					{
						resultLength = inflater.inflate(decompressed);
					}
					catch (DataFormatException e)
					{ throw new IOException(e); }
					
					// undo predictor
					for (int j = 1; j < resultLength; j++)
					{
						int value = (decompressed[j-1] & 0xff)+(decompressed[j] & 0xff)-128;
						decompressed[j] = (byte) value;
					}
					
					// reorder the pixel data
					for (int j = 0; j < resultLength/2; j++)
					{
						pixeldata[2*j+0] = decompressed[j];
						pixeldata[2*j+1] = decompressed[j+resultLength/2];
					}
//					inflater.end();
					return resultLength;
				}
			};
	else if (compression == CompressionType.RLE_COMPRESSION)
		decompressor = new Decompressor()
			{
				private final byte[] decompressed = new byte[maxBlockSize];
				@Override
				public int decompress(byte[] compressed, int size, byte[] pixeldata) throws IOException
				{
					int pos = 0;
					int written = 0;
					while (pos < size)
					{
						if (compressed[pos] < 0)
						{
							int count = -compressed[pos++];
							while (count-- > 0)
								decompressed[written++] = compressed[pos++];
						}
						else
						{
							int count = compressed[pos++]+1;
							if (written+count > pixeldata.length)
								throw new IOException((written+count)+" > "+pixeldata.length);
							while (count-- > 0)
								decompressed[written++] = compressed[pos];
							pos++;
						}
					}
					
					// undo predictor
					for (int j = 1; j < written; j++)
					{
						int value = (decompressed[j-1] & 0xff)+(decompressed[j] & 0xff)-128;
						decompressed[j] = (byte) value;
					}
					
					// reorder the pixel data
					for (int j = 0; j < written/2; j++)
					{
						pixeldata[2*j+0] = decompressed[j];
						pixeldata[2*j+1] = decompressed[j+written/2];
					}
					return written;
				}
			};
	else
		throw new IOException("Not implemented!");
	
	byte[] compressed = new byte[maxBlockSize];
	byte[] pixeldata = new byte[maxBlockSize];
	Buffer buffer = new Buffer(Buffer.LITTLE_ENDIAN, pixeldata);
	
	HdrImage2D image = new HdrImage2D(width, height);
	
	for (int i = 0; i < numBlocks; i++)
	{
		int ystart = in.readInt()-dw.getTop();
		int dataSize = in.readInt();
		if (dataSize > maxBlockSize) throw new IOException("Argh in line "+i);
//		System.out.print("  Block "+i+" starts at "+ystart+" len "+dataSize);
		in.readFully(compressed, 0, dataSize);
		
		int ystop = ystart+blockSize;
		ystop = ystop < height ? ystop : height;
		int expectedSize = (ystop-ystart)*channels.getLineLength(width);
		
		// decompress
		if (dataSize < expectedSize)
		{
			int resultLength = decompressor.decompress(compressed, dataSize, pixeldata);
			if (resultLength != expectedSize)
				throw new IOException("Decompression result on line "+ystart+" does not match expected length. "+
						"I got "+resultLength+" bytes, but expected "+expectedSize+"!");
		}
		else
		{
			System.arraycopy(compressed, 0, pixeldata, 0, dataSize);
		}
		
		// decode pixel data into image
		if (pixelType == PixelType.HALF)
		{
			if (numChannels == 3)
			{
				for (int y = ystart; y < ystop; y++)
				{
					int chsize = width*2;
					int offset = (y-ystart)*channels.getLineLength(width);
					for (int x = 0; x < width; x++)
					{
						float b = buffer.readHalf(offset+0*chsize+2*x);
						float g = buffer.readHalf(offset+1*chsize+2*x);
						float r = buffer.readHalf(offset+2*chsize+2*x);
						image.setRGB(x, y, r, g, b);
					}
				}
			}
			else if (numChannels == 4)
			{
				for (int y = ystart; y < ystop; y++)
				{
					int offset = (y-ystart)*channels.getLineLength(width);
					for (int x = 0; x < width; x++)
					{
						float a = buffer.readHalf(offset+channels.getOffset(0, width)+2*x);
						float b = buffer.readHalf(offset+channels.getOffset(1, width)+2*x);
						float g = buffer.readHalf(offset+channels.getOffset(2, width)+2*x);
						float r = buffer.readHalf(offset+channels.getOffset(3, width)+2*x);
						image.setRGBA(x, y, r, g, b, a);
					}
				}
			}
			else
				throw new RuntimeException("Not implemented!");
		}
		else if (pixelType == PixelType.FLOAT)
		{
			if (numChannels == 3)
			{
				for (int y = ystart; y < ystop; y++)
				{
					int chsize = width*4;
					int offset = (y-ystart)*channels.getLineLength(width);
					for (int x = 0; x < width; x++)
					{
						float b = buffer.readFloat(offset+0*chsize+4*x);
						float g = buffer.readFloat(offset+1*chsize+4*x);
						float r = buffer.readFloat(offset+2*chsize+4*x);
						image.setRGB(x, y, r, g, b);
					}
				}
			}
			else if (numChannels == 4)
			{
				for (int y = ystart; y < ystop; y++)
				{
					int offset = (y-ystart)*channels.getLineLength(width);
					for (int x = 0; x < width; x++)
					{
						float a = buffer.readFloat(offset+channels.getOffset(0, width)+4*x);
						float b = buffer.readFloat(offset+channels.getOffset(1, width)+4*x);
						float g = buffer.readFloat(offset+channels.getOffset(2, width)+4*x);
						float r = buffer.readFloat(offset+channels.getOffset(3, width)+4*x);
						image.setRGBA(x, y, r, g, b, a);
					}
				}
			}
			else
				throw new RuntimeException("Not implemented: "+numChannels);
		}
		else
			throw new IOException("Not implemented!");
	}
	
	return image;
}

}
